from pathlib import Path

from gelslimpi.core import InputSource, GelSlimInterface
from gelslimpi.utils import VirableRateVideoWriter
from gelslimpi import qtcv
from gelslimpi.marker_tracker import *

import numpy as np
import cv2

BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "runs/23_07_08_17_09_53"


def draw_dots(img, dots):
    """绘制点
    :param img: 输入图片
    """
    circle_color = (3, 253, 253) 
    for i, row in enumerate(dots):  # 丢弃最外围点
        for j, dot in enumerate(row):
            cv2.circle(img, dot, 2 , circle_color, 2)  # 黄点
    return img

if __name__ == '__main__':
    data = np.load(DATA_DIR/"data.npz")
    dots = data["dots"]
    forces = data["force"]

    for img_path in DATA_DIR.glob("0*.jpg"):
        id = int(img_path.stem)
        img = cv2.imread(str(img_path))
        img = draw_dots(img, dots[id])
        qtcv.imshow("show", img)
        key = qtcv.waitKey(1000)
        if key == qtcv.Key.Key_Escape:
            break

